import sqlite3
import uuid

# 数据库表名常量
TABLE_M3U8_TASK = "M3u8Task2"
TABLE_TS_BEAN = "TsBean2"

class TaskDB:
    def __init__(self, db_name):
        self.db_name = db_name
        self.conn = self.create_connection()
        self.init_db()
        
    def create_connection(self):
        """创建数据库连接"""
        return sqlite3.connect(self.db_name, check_same_thread=False)
    
    def close_connection(self):
        """关闭数据库连接"""
        if self.conn:
            self.conn.close()
    
    def execute(self, sql, params=()):
        """执行SQL语句"""
        cursor = self.conn.cursor()
        try:
            cursor.execute(sql, params)
            self.conn.commit()
            return cursor
        except sqlite3.Error as e:
            self.conn.rollback()
            raise e
    
    def generate_uid(self):
        """生成唯一ID"""
        return str(uuid.uuid4())
    
    def init_db(self):
        """初始化数据库表结构"""
        sqls = [
            f'''
            CREATE TABLE IF NOT EXISTS {TABLE_M3U8_TASK} (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                url TEXT NOT NULL,
                refer TEXT,
                keyUrl TEXT,
                name TEXT NOT NULL,
                ckey TEXT,
                uid TEXT NOT NULL UNIQUE,
                duration REAL DEFAULT 0,
                state INTEGER DEFAULT 0,
                totalNum INTEGER DEFAULT 0,
                downloaded INTEGER DEFAULT 0,
                save_path TEXT,
                create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
            ''',
            f'''
            CREATE TABLE IF NOT EXISTS {TABLE_TS_BEAN} (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                order_num INTEGER NOT NULL,
                link TEXT NOT NULL,
                second REAL,
                m3u8Id TEXT NOT NULL,
                state INTEGER DEFAULT 0,
                FOREIGN KEY(m3u8Id) REFERENCES {TABLE_M3U8_TASK}(uid)
            )
            '''
        ]
        
        for sql in sqls:
            self.execute(sql)
    
    def create_m3u8_task(self, **kwargs):
        """创建M3U8任务"""
        sql = f'''
        INSERT INTO {TABLE_M3U8_TASK} 
        (url, refer, keyUrl, name, uid, totalNum, save_path) 
        VALUES (?, ?, ?, ?, ?, ?, ?)
        '''
        params = (
            kwargs['url'],
            kwargs['refer'],
            kwargs['keyUrl'],
            kwargs['name'],
            kwargs['uid'],
            kwargs['totalNum'],
            kwargs.get('save_path', '')
        )
        cursor = self.execute(sql, params)
        return cursor.lastrowid
    
    def create_ts_segment(self, **kwargs):
        """创建TS片段任务"""
        sql = f'''
        INSERT INTO {TABLE_TS_BEAN} 
        (order_num, link, second, m3u8Id) 
        VALUES (?, ?, ?, ?)
        '''
        params = (
            kwargs['order_num'],
            kwargs['link'],
            kwargs.get('second', 0),
            kwargs['m3u8Id']
        )
        self.execute(sql, params)
    
    def get_task(self, task_id):
        """获取指定任务信息"""
        sql = f"SELECT * FROM {TABLE_M3U8_TASK} WHERE id=?"
        cursor = self.execute(sql, (task_id,))
        row = cursor.fetchone()
        if not row:
            return None
        
        columns = [column[0] for column in cursor.description]
        return dict(zip(columns, row))
    
    def get_all_tasks(self):
        """获取所有任务"""
        sql = f"SELECT id, name, state, totalNum, downloaded FROM {TABLE_M3U8_TASK} ORDER BY create_time DESC"
        cursor = self.execute(sql)
        tasks = []
        columns = [column[0] for column in cursor.description]
        
        for row in cursor.fetchall():
            tasks.append(dict(zip(columns, row)))
        
        return tasks
    
    def get_ts_segments(self, task_uid):
        """获取指定任务的所有TS片段"""
        sql = f"SELECT * FROM {TABLE_TS_BEAN} WHERE m3u8Id=? AND state=0 ORDER BY order_num"
        cursor = self.execute(sql, (task_uid,))
        segments = []
        columns = [column[0] for column in cursor.description]
        
        for row in cursor.fetchall():
            segments.append(dict(zip(columns, row)))
        
        return segments
    
    def update_ts_state(self, segment_id, state):
        """更新TS片段状态"""
        self.execute(f"UPDATE {TABLE_TS_BEAN} SET state=? WHERE id=?", (state, segment_id))
        
        # 更新任务的已下载数量
        self.execute(f"""
            UPDATE {TABLE_M3U8_TASK}
            SET downloaded = downloaded + 1
            WHERE uid = (
                SELECT m3u8Id FROM {TABLE_TS_BEAN} WHERE id=?
            )
        """, (segment_id,))
    
    def check_task_completed(self, task_uid):
        """检查任务是否完成"""
        sql = f"""
        SELECT COUNT(*) FROM {TABLE_TS_BEAN} 
        WHERE m3u8Id=? AND state=0
        """
        cursor = self.execute(sql, (task_uid,))
        remaining = cursor.fetchone()[0]
        return remaining == 0
    
    def mark_task_completed(self, task_id):
        """标记任务为完成状态"""
        self.execute(f"UPDATE {TABLE_M3U8_TASK} SET state=1 WHERE id=?", (task_id,))
        
    def delete_task(self, task_uid):
        """删除任务及其所有TS片段"""
        # 先删除TS片段
        self.execute(f"DELETE FROM {TABLE_TS_BEAN} WHERE m3u8Id=?", (task_uid,))
        
        # 再删除主任务
        self.execute(f"DELETE FROM {TABLE_M3U8_TASK} WHERE uid=?", (task_uid,))
        
    def close(self):
        """关闭数据库连接"""
        self.close_connection()